import os
import shutil

import ai.AiConstant
from ai.audio_dataset import common_data_handler
from ai.config.config import GxlNode
from ai.utils import utils_file


def handle_sanfangyan_data():
    config = GxlNode.get_config_from_yaml('./config_三方言.yaml')
    output_dir = config.output_dir
    source_dir = config.source_dir
    for prefix in ['train', 'test', 'dev']:
        wav_scp = source_dir + f"{prefix}_wav.scp"
        text_scp = source_dir + f"{prefix}_text"
        output_dir = utils_file.join_path(output_dir, prefix)
        logger = ai.AiConstant.AI_LOGGER('./output/make_shard_to_data_三方言.log')
        common_data_handler.make_shard_file(wav_scp, text_scp, output_dir, num_utt_per_shard=1000,
                                            prefix_for_tar_file=prefix, logger=logger)


def handle_big_dir_data():
    """
    保证每个子目录中均有一个wav.csp文件和text文件， 否则跳过对该目录的处理
    """
    config = GxlNode.get_config_from_yaml('./config_big.yaml')
    output_dir_source = config.output_dir
    source_dir = config.source_dir
    utils_file.makedir_sil(output_dir_source)
    logger = ai.AiConstant.AI_LOGGER('./output/make_shard_to_data_big.log')
    for root, dirs, files in os.walk(source_dir):
        for dir in dirs:
            logger.info(f'开始处理{dir}目录')
            wav_scp = utils_file.join_path(root, dir, "wav.scp")
            text_scp = utils_file.join_path(root, dir, "text")
            if not os.path.exists(wav_scp) or not os.path.exists(text_scp):
                logger.info(f'skip {dir}, the dir do not have wav.scp or text file')
                continue
            if dir == "WenetSpeech":
                logger.info(f'skip {dir}, the dir is Wenetspeech, and the dir has handled before')
                continue
            output_dir = utils_file.join_path(output_dir_source, dir)
            size = common_data_handler.get_dir_size(output_dir)
            if size > 10:
                logger.info(f'skip {dir}, size:{size} the dir is big, and the dir has handled before')
                continue
            if os.path.exists(output_dir):
                shutil.rmtree(output_dir)
            common_data_handler.make_shard_file(wav_scp, text_scp, output_dir, num_utt_per_shard=1000,
                                                prefix_for_tar_file=dir, logger=logger)
        break


def handle_test():
    config = GxlNode.get_config_from_yaml('./config_test.yaml')
    output_dir = config.output_dir
    source_dir = config.source_dir
    wav_scp = source_dir + "wav.scp"
    text_scp = source_dir + "text"
    logger = ai.AiConstant.AI_LOGGER('./output/make_shard_to_data_test.log')
    common_data_handler.make_shard_file(wav_scp, text_scp, output_dir, num_utt_per_shard=1000,
                                        prefix_for_tar_file='yuzhishoujiluyin', logger=logger)


if __name__ == '__main__':
    """
    将数据打成shard, 每个shard的大小为1000utt, 每个文件夹打一个
    """
    # handle_sanfangyan_data()
    handle_big_dir_data()
    # handle_test()
